import torch
from torch.utils.data import DataLoader, WeightedRandomSampler, random_split

from SSQDataset import SSQDataset
from model import SSQNet
from ssqdata import UpdateData
from tools import ReadConfigs, PrintLog
from train import train


def weight_calculate(dataset):
    num1 = 0
    num2 = 0
    list1 = torch.tensor([0, 1], dtype=torch.float32)
    list2 = torch.tensor([1, 0], dtype=torch.float32)
    for data_list, target_list in dataset:
        if torch.equal(target_list, list1):
            num1 += 1
        elif torch.equal(target_list, list2):
            num2 += 1

    weight1 = 1. / num1.__float__()
    weight2 = 1. / num2.__float__()

    weight = []
    for data_list, target_list in dataset:
        if torch.equal(target_list, list1):
            weight.append(weight1)
        elif torch.equal(target_list, list2):
            weight.append(weight2)
    return weight, num1 / (num1 + num2)


def main(index=1):
    # 更新数据
    # UpdateData()

    # 取得yaml中的参数
    configs = ReadConfigs()
    red_file_path = configs["RED_FILE_PATH"]
    sample_nums = configs["SAMPLE_NUMS"]

    # 获取数据集
    ssq_dataset = SSQDataset(red_file_path, index, sample_nums)

    # 获取训练集和数据集的大小
    train_size = int(len(ssq_dataset) * 0.8)
    test_size = len(ssq_dataset) - train_size

    # 分出训练集和数据集
    train_dataset, test_dataset = random_split(ssq_dataset, [train_size, test_size])

    # 训练集
    train_samples_weight, _ = weight_calculate(train_dataset)
    train_sampler = WeightedRandomSampler(train_samples_weight, len(train_samples_weight), replacement=True)
    train_loader = DataLoader(train_dataset, batch_size=int(configs["BATCH_SIZE"]), num_workers=1,
                              sampler=train_sampler)

    # 测试集
    _, down_acc = weight_calculate(test_dataset)
    test_loader = DataLoader(test_dataset, batch_size=int(configs["BATCH_SIZE"]), shuffle=False, num_workers=1)
    PrintLog("DataLoader is ready.")

    # Set the device
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    # Get the model
    model = SSQNet(input_size=33, hidden_size=128, num_layers=3, num_classes=2, sample_nums=sample_nums)
    model = model.to(device)
    print("Model is ready.")

    # Set the optimizer
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.Adam(params, lr=configs["LEARNING_RATE"])

    # epochs
    epochs = int(configs["EPOCHS"])

    train(
        model=model,
        epochs=epochs,
        train_loader=train_loader,
        test_loader=test_loader,
        optimizer=optimizer,
        index=index
    )
    PrintLog("Finished training.")


if __name__ == "__main__":
    for i in range(1, 34):
        main(i)
